﻿//uncomment only one of these
#define normalsend         //producer consumers using MPI send/receive
//#define marshalcomplex   //producer consumers marsalling complex rows as two separate double rows
//#define scattergather    //scatter gather using MPI scatter/gather

//uncomment the following to use immediatesend/receives in scatter gather
//#define immediatesends

//uncomment if verbose progress meesages to console wanted
//#define progressmessages

//uncomment to workaround MPI bugs
#define fixmpi

using System;
using System.Collections.Generic;
using System.Text;
using System.IO;
using System.Collections;
using MPI;
using System.Threading;
using System.Threading.Tasks;
using System.Numerics;

namespace Dwarf.Managed
{

    /// <summary>
    /// Solve the Spectral method problem.
    /// </summary>
    /// <remarks> 
    /// For the demonstration purposes of the project,
    /// we have chosen two-dimensional iterative FFT algorithm
    /// See Solve() below
    /// </remarks>
    class Solver
    {
        /// <summary>
        /// Size of matrix.
        /// </summary>
        public int length;

        /// <summary>
        /// Working array.
        /// </summary>
        public Complex[][] complexArray;

        /// <summary>
        /// Common MPI objects, values, and tests
        /// </summary>
        static public MPI.Intracommunicator MpiWorld { get { return Communicator.world; } }
        static public int MpiRank { get { return MpiWorld.Rank; } }
        static public int MpiRoot { get { return Settings.ROOT_PROCESS; } }
        static public int MpiSize { get { return MpiWorld.Size; } }
        static public bool IsMpiRoot { get { return MpiRank == MpiRoot; } }

#if scattergather
#if immediatesends
        private MPI.ReceiveRequest[][] gatherRecvReqs;
        private MPI.ReceiveRequest[] scatterRecvReqs;
#endif
#else  
        // normalsend or marshalcomplex
        /// <summary>
        /// Row index that indicates all done to workers
        /// </summary>
        static private int AllDoneTag;
#endif
#if marshalcomplex
        /// <summary>
        /// marshalling send/recv buffers
        /// </summary>
        private double[] recvBuf;
        private double[] sendBuf;
#endif


        /// <summary>
        /// Default constructor
        /// </summary>
        public Solver() 
        {
        }

        //Solve using overlapped producer consumer MPI

        /// <summary>
        /// Transposition of matrix using TPL
        /// </summary>
        /// <remarks>this routine is so fast that TPL or not TPL is not noticeable different</remarks>
        static private void transpose(Complex[][] matrix) {
            var length = matrix.Length;
            Parallel.For(0, length - 1, i =>
            {
                //for (int i = p; i < p + step; i++)
                //{
                for (int j = i + 1; j < length; j++) {
                    var temp = matrix[i][j];
                    matrix[i][j] = matrix[j][i];
                    matrix[j][i] = temp;
                }
                //}
            });
        }

        private Complex[] W;

        /// <summary>
        /// 1-dimension Fast Fourier Algorithm
        /// </summary>
        private void FFT(Complex[] complexLine)
        {         
            if (complexLine == null) return;

            //variables for the FFT
            var length = complexLine.Length;
            int bitwidth = (int)Math.Log(length, 2);

            if (W == null || W.Length != length) {
                //compute phase weights
                W = new Complex[length];

                for (var i = 0; i < length; i++) {
                    W[i] = new Complex(
                        Math.Cos(-((2 * i * Math.PI) / length)),
                        Math.Sin(-((2 * i * Math.PI) / length)));
                }
            }

            //reorder for butterflies
            int firstBit = length >> 1;
            int index = 0;
            int reverseIndex = 0;
            for (; ; ) {
                if (reverseIndex < index) {
                    var w = complexLine[index];
                    complexLine[index] = complexLine[reverseIndex];
                    complexLine[reverseIndex] = w;
                }
                index++;
                if (index == length)
                    // Break here, because the code below doesn't terminate for all 1's.
                    break;

                // Add 1 to reverse index, going from left to right:
                int carry = firstBit;
                // If adding 1 changed the bit to 0, we have a carry:
                while (((reverseIndex ^= carry) & carry) == 0)
                    carry >>= 1;
            }

            //combine and untangle applying phase weights
            //using successive butterflies of 2,4,...,2**bitwidth 
            for (var i = 1; i <= bitwidth; i++)
            {
                var halfStep = (1 << (i - 1));
                for (var j = 0; j < halfStep; j++)
                {
                    var tempShift = (j << (bitwidth - i));
                    var w = W[tempShift];
                    for (var k = j; k < length; k += halfStep * 2)
                    {
                        var u = complexLine[k + halfStep];
                        var v = complexLine[k];
                        var t = w * u;
                        complexLine[k + halfStep] = v - t;
                        complexLine[k] = v + t;
                    }
                }
            }
        }

        /// <summary>
        /// Based method.
        /// </summary>
        public void Solve()
        {
#if scattergather

            //rows and indexs of rows being processed locally
            Complex[][] rows = null;
            int [] indexs = null;

            /*Scatter rows to all MPI processes */
            ScatterRows(complexArray, length, ref rows, ref indexs);

            /* Transform the rows */
            Parallel.ForEach(rows, row => FFT(row));

            /* Gather transfomed rows on root process */
            GatherRows(complexArray, rows, indexs);

            /*Transposition */
            if (IsMpiRoot) transpose(complexArray);

            /*Scatter rows to all MPI processes */
            ScatterRows(complexArray, length, ref rows, ref indexs);

            /*Transform the columns(after transposition) */
            Parallel.ForEach(rows, row => FFT(row));

            /* Gather transfomed rows on root process */
            GatherRows(complexArray, rows, indexs);

            /*Transposition */
            if (IsMpiRoot) transpose(complexArray);

#else //not scatter gather  -- sendnormal or complexmarshal

            //transform rows
            EvaluateRowsUsingMPI(complexArray, FFT);

            /*Transposition to switch rows and columsn*/
            if (IsMpiRoot) transpose(complexArray);

            //transform rows (columns prior to trnspose)
            EvaluateRowsUsingMPI(complexArray, FFT);

            /*Transposition to restore rows and columns*/
            if (IsMpiRoot) transpose(complexArray);
#endif
        }

#if scattergather

        /// <summary>
        /// Scatters the rows of matrix to all MPI processes evenly
        /// </summary>
        /// <remarks>
        /// Matrix is cleared of all rows on return
        /// </remarks>
        /// <typeparam name="T">Element type of the matrix</typeparam>
        /// <param name="matrix">Matrix to scatter</param>
        /// <param name="numRows">Number of rows of matrix to scatter</param>
        /// <param name="rows">Set on return to the rows scattered to the caller</param>
        /// <param name="indexs">Set on return to the corresponding row indices for each row in rows</param>
        void ScatterRows(Complex[][] matrix, int numRows, ref Complex[][] rows, ref int[] indexs)
        {
            //compute how to carve up matrix 

            var procNumRows = (numRows + (MpiSize - 1)) / MpiSize;
#if fixmpi //Don't use MPI scatter, emulate it ourselves with send/receive

            //make sure we have the outputs allocated
            rows = rows ?? new Complex[procNumRows][];
            indexs = indexs ?? new int[procNumRows];

            //root has the data to scatter
            if (IsMpiRoot)
            {
                //return a batch of rows to each
                var index = 0;

                //place roots portion in rows and indexs
                for (var i = 0; i < procNumRows; i++)
                {
                    rows[i] = (i < matrix.Length)
                        ? matrix[i]
                        : new Complex[length];
                    indexs[i] = i;
                }
                index = procNumRows;

                //send rest to workers
                for (var i = 1; i < MpiSize; i++)
                {
                    for (var j = 0; j < procNumRows; j++)
                    {
                        var row = (index < matrix.Length)
                            ? matrix[index]
                            : new Complex[length];
#if immediatesends
                        MpiWorld.ImmediateSend(row, i, index);
#else
                        MpiWorld.Send(row, i, index);
#endif
                        index++;
                    }
                }
            }
            else
            {
                //workers receive their portions
#if immediatesends
                //Fire off all requests
                scatterRecvReqs = scatterRecvReqs ?? new MPI.ReceiveRequest[procNumRows];
                var first = MpiRank * procNumRows;
                for (var i = 0; i < procNumRows; i++)
                {
                    rows[i] = rows[i] ?? new Complex[length];
                    scatterRecvReqs[i] = MpiWorld.ImmediateReceive(MpiRoot, first + i, rows[i]);
                    indexs[i] = first + i;
                }
                //make sure the data arrives
                foreach (var req in scatterRecvReqs) req.Wait();
#else // not immediatesends
                MPI.Status probe = null;
                for (var i = 0; i < procNumRows; i++)
                {
                    probe = MpiWorld.Probe(MpiRoot, Communicator.anyTag);
                    rows[i] = (rows[i] == null) ? new Complex[length] : rows[i];
                    MpiWorld.Receive(probe.Source, probe.Tag, ref rows[i]);
                    indexs[i] = probe.Tag;
                }
#endif //immediatesends
            }


#else  //not fixmpi --- Use MPI scatter gather      
            //carve the array up on the root
            var procRows = new Complex[MpiSize][][];
            var procIndexs = new int[MpiSize][];
            if (IsMpiRoot)
            {
                // group rows for distribution to the processors
                var index = 0;
                for (var i = 0; i < MpiSize; i++)
                {
                    procRows[i] = new Complex[procNumRows][];
                    procIndexs[i] = new int[procNumRows];
                    for (var j = 0; j < procNumRows; j++)
                    {
                        procIndexs[i][j] = index;
                        if (index < matrix.Length)
                        {
                            //move row to the group of rows going to process i 
                            procRows[i][j] = matrix[index];
                            index++;
                        }
                    }
                }
            }  //IsMpiRoot

            //distribute using MPI and return local rows and indexs received
           //BUG: MPI won't scatter int[][] procIndexs
            indexs = Communicator.world.Scatter(procIndexs, Settings.ROOT_PROCESS);

            //BUG: MPI won't scatter Complex[][][] ---  procRows
            rows = Communicator.world.Scatter(procRows, Settings.ROOT_PROCESS);
#endif //fixmpi

#if progressmessages
            Console.WriteLine("Process {0} GOT {1} rows of size {2} from scatter",
                MpiRank, rows.Length, rows[0].Length);
#endif
        }

        /// <summary>
        /// Gather rows from all MPI rocesses and reconstruct matrix on root
        /// </summary>
        /// <remarks>
        /// matix is only meaningful on the root process.
        /// </remarks>
        /// <typeparam name="T">Element type of matrix</typeparam>
        /// <param name="matrix">matrix to reconstruct</param>
        /// <param name="rows">local rows to this process</param>
        /// <param name="rowIndexs">Corresponding indices for each row</param>
        void GatherRows(Complex[][] matrix, Complex[][] rows, int[] rowIndexs)
        {
            var procNumRows = rows.Length;

#if fixmpi //BUG:??not sure?? MPI wont gather int[]s and/or Complex[][]s
            
            //Do it with send receive instead
            if (!IsMpiRoot)
            {
                //Worker process just send rows with index as tag
                for (var i = 0; i < procNumRows; i++)
                {
#if immediatesends
                    MpiWorld.ImmediateSend(rows[i], MpiRoot, rowIndexs[i]);
#else
                    MpiWorld.Send(rows[i], MpiRoot, rowIndexs[i]); 
#endif
                }
#if progressmessages
                Console.WriteLine("Process {0} gather {1} rows to Root", 
                    MpiRank, rows.Length);
#endif
            }
            else
            {
                //Root grabs tagged rows and shuffles them back into matrix
                //First shuffle rows already in root
                for (var i = 0; i < procNumRows; i++)
                { matrix[rowIndexs[i]] = rows[i]; }
                var processed = procNumRows;

                //Now receive tagged rows from workers and shuffle in
#if immediatesends
                //fire off all of the requests
                gatherRecvReqs = gatherRecvReqs ?? new MPI.ReceiveRequest[MpiSize-1][];
                for (var n = 1; n < MpiSize; n++)
                {
                    var procReqs = gatherRecvReqs[n - 1] ?? new MPI.ReceiveRequest[procNumRows];
                    gatherRecvReqs[n - 1] = procReqs;
                    for (var i = 0; i < procNumRows; i++)
                    {
                        if (processed < matrix.Length)
                        {
                            complexArray[i] = matrix[i] ?? new Complex[rows[0].Length];
                            procReqs[i] = MpiWorld.ImmediateReceive(n, processed, complexArray[processed]);
                            processed++;
                        }
                    }
                }
                //wait until the data arrives
                foreach (var reqs in gatherRecvReqs) { foreach (var req in reqs) { req.Wait(); } };

#else //not immediatesends
                Complex[] emptyRow = null;
                while (processed < matrix.Length)
                {
                    var probe = MpiWorld.Probe(Communicator.anySource, Communicator.anyTag);
                    if (probe.Tag < matrix.Length)
                    {
                        MpiWorld.Receive(probe.Source, probe.Tag, ref matrix[probe.Tag]);
                        processed++;
                    }
                    else
                    {
                        emptyRow = emptyRow ?? new Complex[matrix[0].Length];
                        MpiWorld.Receive(probe.Source, probe.Tag, ref emptyRow);
                    }
                }
#endif //immediatesends

#if progressmessages
                Console.WriteLine("Root shuffled all rows into matrix}");
#endif
            }

#else  //not fixmpi ---  using MPI Gather
            var indexs = Communicator.world.Gather(rowIndexs, Settings.ROOT_PROCESS);
            var values = Communicator.world.Gather(rows, Settings.ROOT_PROCESS);

            //only root gets non-null
            if (values != null && indexs != null)
            {
                //shuffle rows back into matrix
                for (var i = 0; i < values.Length; i++)
                {
                    var procRows = values[i];
                    var procIndex = indexs[i];
#if progressmessages
                    Console.WriteLine("Process {0} group {1} has {2} rows and {3} indices",
                        Communicator.world.Rank, i, procRows.Length, procIndex.Length);
#endif
                    for (var j = 0; j < procRows.Length; j++)
                    {
                        if (procIndex[j] != no_index)
                        { matrix[procIndex[j]] = procRows[j]; }
                    }
                }
            }
#endif //fixmpi
        }

#endif //scattergather

#if scattergather
#else //not scattergather -- normalsend or marshalcomplex

        /// <summary>
        /// Evaluate an action on every row in matrix
        /// </summary>
        /// <remarks>
        /// The action will be evaluated once for each row 
        /// with the row passed to the action as the sole argument, 
        /// The action is allowed to change the contents of the row
        /// </remarks>
        /// <typeparam name="T">Element type of matrix</typeparam>
        /// <param name="matrix">matrix to operate on</param>
        /// <param name="Row_Action">Single argument action to perform</param>
        private void EvaluateRowsUsingMPI(Complex[][] matrix, Action<Complex[]> action)
        {
            //set bounds of matrix for worker processes
            var colLength = length;
            var rowLength = length;

            //Set All Done tag for worker processes
            AllDoneTag = colLength;

            if (IsMpiRoot)
            {
                //Root Process
                var rowsSent = 0;
                var rowsProcessed = 0;

                //Give each worker process some rows to start with
                const int batch = 1;  //number of rows to give each worker at startup
                for (var j = 0; j < batch; j++)
                {
                    for (var i = 1; i < MpiSize; i++)
                    {
                        SendRow(matrix[rowsSent % colLength], i, rowsSent);
                        rowsSent++;
                    }
                }

                //Loop until all rows sent and all transformed rows received
                while (rowsProcessed < colLength)
                {
                    //Check if any results sent back from any worker
                    var probeStatus = MpiWorld.ImmediateProbe(Communicator.anySource, Communicator.anyTag);
                    if (probeStatus == null)
                    {
                        //nothing to receive, do we have any rows left to do
                        if (rowsSent < colLength)
                        {
                            //yes, transform next row in root process
                            action(matrix[rowsSent]);
                            rowsSent++;
                            rowsProcessed++;
                        }
                        else { Thread.Sleep(0); }  //we have nothing to do
                        continue;
                    }

                    //receive row, waiting if necessary                
                    var recvStatus = ReceiveRow(probeStatus.Source, ref matrix[probeStatus.Tag], probeStatus.Tag);
                    var source = recvStatus.Source;
                    var recvIndex = recvStatus.Tag;
                    rowsProcessed++;

                    //Send next row (if any) to process  that completed.
                    if (rowsSent < colLength)
                    {
                        SendRow(matrix[rowsSent], source, rowsSent);
                        rowsSent++;
                    }
                }

                //Tell workers all done
                var row = new Complex[rowLength];
                for (var i = 1; i < MpiSize; i++)
                {
                    SendRow(row, i, AllDoneTag);
                }
            }
            else
            {
                //Worker Processes

                //Process rows until told to stop.
                var row = new Complex[rowLength];   
                while (true)
                {
                    //Receive row, waiting if necessary.
                    int rowIndex;
                    var recvStatus = ReceiveRow(MpiRoot, ref row, Communicator.anyTag);
                    rowIndex = recvStatus.Tag;

                    //are we done?
                    if (rowIndex == AllDoneTag) break;

                    /*Transform the row */
                    action(row);

                    /* send the transformed row back */
                    SendRow(row, MpiRoot, rowIndex);

                }
            }
        }

        /// <summary>
        /// Send row and index of row to dest
        /// </summary>
        /// <param name="tRow">array of type Complex[] to send</param>
        /// <param name="dest">MPI rank of process to receive row</param>
        /// <param name="index">Index of row in original matrix</param>         
        private void SendRow(Complex[] row, int dest, int index)

#if normalsend
        { 
            MpiWorld.Send(row, dest, index);
        }
#elif marshalcomplex
        {
            //send reals
            if (sendBuf == null) sendBuf = new double[row.Length];
            for (var i = 0; i < row.Length; i++) sendBuf[i] = row[i].real;
            MpiWorld.Send(sendBuf, dest, index);

            //send imags
            for (var i = 0; i < row.Length; i++) sendBuf[i] = row[i].imag;
            MpiWorld.Send(sendBuf, dest, index);
#if progressmessages
            Console.WriteLine("Process {0} SENT row {1} to {2} size {3}",
                MpiRank, index, dest, row.Length);
#endif
        }
#endif


        /// <summary>
        /// Receive a row from requested source
        /// </summary>
        /// <param name="source">MPI rank to receive row from.  To receive from any, Use Communicator.anySource</param>
        /// <param name="row">Array to receive the row</param>
        /// <param name="index">Set on return to the row index of the received row in the orginalmatrix</param>
        /// <returns>MPI.CompletedStatus of received row</returns>
        private CompletedStatus ReceiveRow(int source, ref Complex[] row, int index)
        {
            CompletedStatus recvStatus;
#if normalsend
            MpiWorld.Receive(source, index, ref row, out recvStatus);
#elif marshalcomplex
            recvBuf = recvBuf ?? new double[row.Length];

            //receive reals
            MpiWorld.Receive(source, index, ref recvBuf, out recvStatus);
            for (var i = 0; i < row.Length; i++) row[i].real = recvBuf[i];

            //receive imags
            MpiWorld.Receive(recvStatus.Source, recvStatus.Tag, ref recvBuf, out recvStatus);
            for (var i = 0; i < row.Length; i++) row[i].imag = recvBuf[i];
#if progressmessages
            Console.WriteLine("Process {0} RECVD row {1} from {2} size {3}", 
                MpiRank, recvStatus.Tag, recvStatus.Source, row.Length);
#endif
#endif //marshalcomplex
            return recvStatus;
        }
#endif //not scattergather
    }
}
